#include "njm_img.hpp"

#define STB_IMAGE_IMPLEMENTATION
#include "thirdParty/stb-master/stb_image.h"

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "thirdParty/stb-master/stb_image_write.h"

namespace njm
{

    NjmImgFormat NjmImg::GetImageFormat(int channels)
    {
        switch (channels)
        {
        case 3:
            return NJM_IMG_FORMAT_RGB8;
        case 4:
            return NJM_IMG_FORMAT_RGBA8;
        case 6:
            return NJM_IMG_FORMAT_RGB16F;
        case 8:
            return NJM_IMG_FORMAT_RGBA16F;
        case 12:
            return NJM_IMG_FORMAT_RGB32F;
        case 16:
            return NJM_IMG_FORMAT_RGBA32F;
        default:
            return NJM_IMG_FORMAT_UNKNOWN;
        }
    }
    int NjmImg::GetChannels(NjmImgFormat image_format)
    {
        switch (image_format)
        {
        case NJM_IMG_FORMAT_RGB8:
            return 3;
        case NJM_IMG_FORMAT_RGBA8:
            return 4;
        case NJM_IMG_FORMAT_RGB16F:
            return 6;
        case NJM_IMG_FORMAT_RGBA16F:
            return 8;
        case NJM_IMG_FORMAT_RGB32F:
            return 12;
        case NJM_IMG_FORMAT_RGBA32F:
            return 16;
        default:
            return 0;
        }
    }

    NjmImgExt NjmImg::GetImageExt(const std::string &path)
    {
        std::string ext = path.substr(path.find_last_of(".") + 1);
        if (ext == "png")
        {
            return NJM_IMG_EXT_PNG;
        }
        else if (ext == "jpg")
        {
            return NJM_IMG_EXT_JPG;
        }
        else if (ext == "bmp")
        {
            return NJM_IMG_EXT_BMP;
        }
        else if (ext == "tga")
        {
            return NJM_IMG_EXT_TGA;
        }
        else
        {
            return NJM_IMG_EXT_UNKNOWN;
        }
    }
    std::string NjmImg::GetImageExtStr(NjmImgExt ext)
    {
        switch (ext)
        {
        case NJM_IMG_EXT_PNG:
            return "png";
        case NJM_IMG_EXT_JPG:
            return "jpg";
        case NJM_IMG_EXT_BMP:
            return "bmp";
        case NJM_IMG_EXT_TGA:
            return "tga";
        default:
            return "";
        }
    }

    VkFormat NjmImg::GetVkFormat(NjmImgFormat image_format)
    {
        switch (image_format)
        {
        case NJM_IMG_FORMAT_RGB8:
            return VK_FORMAT_R8G8B8_UNORM;
        case NJM_IMG_FORMAT_RGBA8:
            return VK_FORMAT_R8G8B8A8_UNORM;
        case NJM_IMG_FORMAT_RGB16F:
            return VK_FORMAT_R16G16B16_SFLOAT;
        case NJM_IMG_FORMAT_RGBA16F:
            return VK_FORMAT_R16G16B16A16_SFLOAT;
        case NJM_IMG_FORMAT_RGB32F:
            return VK_FORMAT_R32G32B32_SFLOAT;
        case NJM_IMG_FORMAT_RGBA32F:
            return VK_FORMAT_R32G32B32A32_SFLOAT;
        case NJM_IMG_FORMAT_UNKNOWN:
            return VK_FORMAT_UNDEFINED;
        default:
            throw std::runtime_error("Unknown NjmImgFormat: " + std::to_string(image_format));
        }
    }

    NjmImgFormat NjmImg::GetNjmImgFormat(VkFormat vk_format)
    {
        switch (vk_format)
        {
        case VK_FORMAT_R8G8B8_UNORM:
            return NJM_IMG_FORMAT_RGB8;
        case VK_FORMAT_R8G8B8_SRGB:
            return NJM_IMG_FORMAT_RGB8;
            // TODO
        case VK_FORMAT_R8G8B8A8_UNORM:
            return NJM_IMG_FORMAT_RGBA8;
        case VK_FORMAT_R8G8B8A8_SRGB:
            return NJM_IMG_FORMAT_RGBA8;
        case VK_FORMAT_R16G16B16A16_SFLOAT:
            return NJM_IMG_FORMAT_RGBA16F;
        case VK_FORMAT_R32G32B32A32_SFLOAT:
            return NJM_IMG_FORMAT_RGBA32F;
        default:
            return NJM_IMG_FORMAT_UNKNOWN;
        }
    }
    //
    NjmImg::NjmImg(const std::string &filepath)
    {
        if (!LoadFromFile(filepath))
        {
            throw std::runtime_error("Failed to load image from file: " + filepath);
        }
    }
    NjmImg::NjmImg(const std::string &name, const std::string &filepath)
    {
        if (!LoadFromFile(name, filepath))
        {
            throw std::runtime_error("Failed to load image from file: " + filepath);
        }
    }

    NjmImg NjmImg::createFromFile(const std::string &filepath)
    {
        NjmImg img;
        if (!img.LoadFromFile(filepath))
        {
            throw std::runtime_error("Failed to load image from file: " + filepath);
        }
        return img;
    }
    NjmImg NjmImg::createFromFile(const std::string &name, const std::string &filepath)
    {
        NjmImg img;
        if (!img.LoadFromFile(name, filepath))
        {
            throw std::runtime_error("Failed to load image from file: " + filepath);
        }
        return img;
    }
    // member load/store
    bool NjmImg::LoadFromFile(const std::string &filepath)
    {
        auto rst = LoadFromFile(filepath, data, width, height, channels);
        if (rst)
        { format = GetImageFormat(channels); }
        return rst;
    }
    bool NjmImg::LoadFromFile(const std::string &name, const std::string &filepath)
    {
        this->name = name;
        auto rst = LoadFromFile(filepath, data, width, height, channels);
        if (rst)
        { format = GetImageFormat(channels); }
        return rst;
    }
    bool NjmImg::SaveToFile(const std::string &filepath)
    {
        return SaveToFile(filepath, data, width, height, channels);
    }

    // static general load/store
    bool NjmImg::LoadFromFile(const std::string &filepath, std::vector<unsigned char> &buffer, int &width, int &height, int &channels)
    {
        NjmImgFormat image_format;
        unsigned char *data = stbi_load(filepath.c_str(), &width, &height, &channels, 0);
        if (data == nullptr)
        {
            std::cerr << "Failed to load image from file: " << filepath << std::endl;
            return false;
        }
        image_format = GetImageFormat(channels);
        if (image_format == NJM_IMG_FORMAT_UNKNOWN)
        {
            std::cerr << "Failed to get image format from channels: " << channels << std::endl;
            return false;
        }
        buffer.resize(width * height * channels);
        memcpy(buffer.data(), data, width * height * channels);
        stbi_image_free(data);
        return true;
    }

    bool NjmImg::SaveToFile(const std::string &filepath, const std::vector<unsigned char> &buffer, int width, int height, int channels)
    {
        int image_format = GetImageFormat(channels);
        if (image_format == NJM_IMG_FORMAT_UNKNOWN)
        {
            std::cerr << "Failed to get image format from channels: " << channels << std::endl;
            return false;
        }
        switch (image_format)
        {
        case NJM_IMG_FORMAT_RGB8:
            stbi_write_png(filepath.c_str(), width, height, channels, buffer.data(), width * channels);
            break;
        case NJM_IMG_FORMAT_RGBA8:
            stbi_write_png(filepath.c_str(), width, height, channels, buffer.data(), width * channels);
            break;
        case NJM_IMG_FORMAT_RGB16F:
            stbi_write_png(filepath.c_str(), width, height, channels, buffer.data(), width * channels * sizeof(float));
            break;
        case NJM_IMG_FORMAT_RGBA16F:
            stbi_write_png(filepath.c_str(), width, height, channels, buffer.data(), width * channels * sizeof(float));
            break;
        case NJM_IMG_FORMAT_RGB32F:
            stbi_write_png(filepath.c_str(), width, height, channels, buffer.data(), width * channels * sizeof(float));
            break;
        case NJM_IMG_FORMAT_RGBA32F:
            stbi_write_png(filepath.c_str(), width, height, channels, buffer.data(), width * channels * sizeof(float));
            break;
        default:
            std::cerr << "Failed to save image to file: " << filepath << std::endl;
            return false;
        }
        return true;
    }

}
